Repository: krayzpipes/txt-ferret

Issue:
Add gzip support
Would be nice to be able to stream in gzip files. Currently gzipped files raise an 'UnicodeDecodeError' because the utf-8 codec can't read the byte in position 1 of the gzipped file.

Comments:

Patch:
diff --git a/README.md b/README.md
index fb7e3d9..bfd38ce 100644
--- a/README.md
+++ b/README.md
@@ -4,7 +4,7 @@ Identify and classify data in your text files with Python.
 ## Description
 **Definition:** txtferret
 - A weasel-like mammal that feasts on rodents... and apparently social security numbers,
-credit card numbers, or any other data that's in your text files.
+credit card numbers, or any other data that's in your text or gzipped text files.
 
 Use custom regular expressions and sanity checks (ex: `luhn` algorithm for account numbers) to find
 sensitive data in virtually any size file via your command line.
@@ -109,12 +109,12 @@ filters:
 - **Pattern:**
     - The regular expression which will be used to find data in the file.
     - Regular expression must be compatible with the python `re` module in the standard library.
-    - Be sure that your regular expression only contains ONE capture group. For example,
+    - Be sure that your regular expression only contains ONE and ONLY ONE capture group. For example,
     if you are capturing a phone number:
         - Don't do this: `'(555-(867|555)-5309)'`
         - Do this: `'(555-(?:867|555)-5309)'`
-        - The first has two capture groups, and inner and an outer.
-        - The second has one capture group (the outer). The inner is a non-capturing
+        - The former example has two capture groups, and inner and an outer.
+        - The latter has one capture group (the outer). The inner is a non-capturing
         group as defined by starting the capture group with `?:`.
     - __Note: If you run into issues with loading a custom filter, try adding
     single-quotes around your regular expression.__
@@ -251,6 +251,10 @@ sanity check which can be paired with a DLP solution. Here are some things it wa
     - No outrageous licensing per GB of data scanned.
 - __You can contribute!__
 
+## Releases
+
+#### Version 0.0.3 - 2019-06-01
+- Added gzip detection and support.
 
 
 # Development
diff --git a/setup.py b/setup.py
index a133a6b..bf4f22a 100644
--- a/setup.py
+++ b/setup.py
@@ -14,7 +14,7 @@ from setuptools import setup, find_packages
 from codecs import open
 from os import path
 
-__version__ = "0.0.2"
+__version__ = "0.0.3"
 description = "Scan text files for sensitive (or non-sensitive) data."
 
 here = path.abspath(path.dirname(__file__))
diff --git a/src/txtferret/_sanity.py b/src/txtferret/_sanity.py
index f591755..0e1452e 100644
--- a/src/txtferret/_sanity.py
+++ b/src/txtferret/_sanity.py
@@ -18,6 +18,10 @@ def luhn(account_string):
         test.
     """
 
+    # TODO - Is there a more effecient way to do this?
+    if not isinstance(account_string, str):
+        account_string = account_string.decode("utf-8")
+
     no_special_chars = re.sub("[\W_]", "", account_string)
 
     try:
diff --git a/src/txtferret/core.py b/src/txtferret/core.py
index 0bf3ea5..102e68a 100644
--- a/src/txtferret/core.py
+++ b/src/txtferret/core.py
@@ -1,6 +1,7 @@
 """Core classes and functions for txt_ferret."""
 
 from datetime import datetime
+import gzip
 from pathlib import Path
 import re
 
@@ -22,10 +23,19 @@ def tokenize(clear_text, mask, index, tokenize=True, show_matches=False):
     :param show_matches: Bool representing whether the clear text should
         be redacted all together or not.
     """
+
     if not show_matches:
         return "REDACTED"
+
+    # byte string can be present if source file is Gzipped
+    # convert to utf-8 string for logging/file output.
+    if not isinstance(clear_text, str):
+        clear_text = clear_text.decode("utf-8")
+        mask = mask.decode("utf-8")
+
     if not tokenize:
         return clear_text
+
     return _get_tokenized_string(clear_text, mask, index)
 
 
@@ -36,6 +46,7 @@ def _get_tokenized_string(text, mask, index):
         than the original string, then the mask will be cut down to
         size.
     """
+
     end_index = index + len(mask)
     text_length = len(text)
     if (text_length - 1) < end_index:
@@ -61,6 +72,30 @@ def _byte_code_to_string(byte_code):
     return bytes((code_,)).decode("utf-8")
 
 
+def gzipped_file_check(file_to_scan, _opener=None):
+    """ Return bool based on if opening file returns UnicodeDecodeError
+
+    If UnicodeDecodeError is returned when trying to read a line of the
+    file, then we will assume this is a gzipped file.
+
+    :param file_to_scan: String containing file path/name to read.
+    :param _opener: Used to pass file handler stub for testing.
+
+    :return: True if UnicodeDecodeError is detected. False if not.
+    """
+
+    # Use test stub or the normal 'open'
+    _open = _opener or open
+
+    try:
+        with _open(file_to_scan, "r") as rf:
+            _ = rf.readline()
+    except UnicodeDecodeError:
+        return True
+    else:
+        return False
+
+
 class Filter:
     """ Helper class to  hold filter configurations and add a simple
     API to interface with Filter attributes.
@@ -75,7 +110,7 @@ class Filter:
         mask should start being applied.
     """
 
-    def __init__(self, filter_dict):
+    def __init__(self, filter_dict, gzip):
         """Initialize the Filter object. Lots handling input from
         the config file here.
 
@@ -109,6 +144,11 @@ class Filter:
                 f"default tokenization mask and index."
             )
 
+        # If gzip, we need to use byte strings instead of utf-8
+        if gzip:
+            self.token_mask = self.token_mask.encode("utf-8")
+            self.pattern = self.pattern.encode("utf-8")
+
         try:
             self.token_index = int(filter_dict["tokenize"].get("index", 0))
         except ValueError:
@@ -124,6 +164,7 @@ class TxtFerret:
     config/settings file or CLI arguments/switches.
 
     :attribute file_name: The name of the file to scan.
+    :attribute gzip: Bool depicting if input file is gzipped
     :attribute tokenize: Determines if txt_ferret will tokenize the
         output of strings that match and pass sanity checks.
     :attribute log_level: Log level to be used by logouru.logger.
@@ -144,9 +185,16 @@ class TxtFerret:
     def __init__(self, file_name=None, config_file=None, config_=None, **cli_settings):
         """Initialize the TxtFerret object."""
         config = config_ or load_config(
-            yaml_file=config_file, default_override=cli_settings["config_override"],
+            yaml_file=config_file, default_override=cli_settings["config_override"]
         )
         self.file_name = file_name
+        self.gzip = gzipped_file_check(self.file_name)
+
+        if self.gzip:
+            logger.info(
+                f"Detected non-text file '{file_name}'... "
+                f"attempting GZIP mode (slower)."
+            )
 
         # Set settings from file.
         self.set_attributes(**config["settings"])
@@ -158,7 +206,9 @@ class TxtFerret:
         self.failed_sanity = 0
         self.passed_sanity = 0
 
-        self.filters = [Filter(filter_dict=filter_) for filter_ in config["filters"]]
+        self.filters = [
+            Filter(filter_dict=filter_, gzip=self.gzip) for filter_ in config["filters"]
+        ]
 
     def set_attributes(self, **kwargs):
         """Sets attributes for the TxtFerret object.
@@ -210,9 +260,17 @@ class TxtFerret:
 
         file_to_scan = file_name or self.file_name
 
-        with open(file_to_scan, "r") as rf:
+        if not self.gzip:
+            _open = open
+        else:
+            _open = gzip.open
+
+        with _open(file_to_scan, "r") as rf:
             for index, line in enumerate(rf):
 
+                # if isinstance(line, bytes):
+                #    line = str(line)
+
                 # If delimiter, then treat file as if it has columns.
                 if self.delimiter:
                     self._scan_delimited_line(line, index)


Test patch:
diff --git a/tests/core/test_core.py b/tests/core/test_core.py
new file mode 100644
index 0000000..8679f1a
--- /dev/null
+++ b/tests/core/test_core.py
@@ -0,0 +1,35 @@
+from contextlib import contextmanager
+
+from txtferret.core import gzipped_file_check
+
+
+def test_gzipped_file_check_return_true():
+
+    @contextmanager
+    def opener_stub_raise_error(x, y):
+
+        class FileHandlerStub:
+
+            @staticmethod
+            def readline():
+                raise UnicodeDecodeError("fake", b"o", 1, 2, "fake")
+
+        yield FileHandlerStub()
+
+    assert gzipped_file_check("f.txt", _opener=opener_stub_raise_error) == True
+
+
+def test_gzipped_file_check_return_false():
+
+    @contextmanager
+    def opener_stub_no_error(x, y):
+
+        class FileHandlerStub:
+
+            @staticmethod
+            def readline():
+                return ""
+
+        yield FileHandlerStub()
+
+    assert gzipped_file_check("f.txt", _opener=opener_stub_no_error) == False