source: file: Add support for .zip file

Author: sudharsana-kjl

CHANGELOG.md

@@ -24,6 +24,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Added support for lzma file source
 - Added support for xz file source
 - Added Data Flow Facilitator
+- Added support for zip file source
 ### Changed
 - Restructured documentation to docs folder and moved from rST to markdown
 - Git feature cloc logs if no binaries are in path

dffml/source/file.py

@@ -6,8 +6,11 @@
 import gzip
 import bz2
 import lzma
+import zipfile
+import io
 from .source import Source
 from .log import LOGGER
+from contextlib import contextmanager
 
 LOGGER = LOGGER.getChild('file')
 
@@ -49,6 +52,8 @@ async def _open(self):
         elif self.filename[::-1].startswith(('.xz')[::-1]) or \
                 self.filename[::-1].startswith(('.lzma')[::-1]):
             opener = lzma.open(self.filename, 'rt')
+        elif self.filename[::-1].startswith(('.zip')[::-1]):
+            opener = self.zip_opener_helper()
         else:
             opener = open(self.filename, 'r')
         with opener as fd:
@@ -66,11 +71,28 @@ async def _close(self):
             elif self.filename[::-1].startswith(('.xz')[::-1]) or \
                     self.filename[::-1].startswith(('.lzma')[::-1]):
                 close = lzma.open(self.filename, 'wt')
+            elif self.filename[::-1].startswith(('.zip')[::-1]):
+                close = self.zip_closer_helper()
             else:
                 close = open(self.filename, 'w')
             with close as fd:
                 await self.dump_fd(fd)
 
+    @contextmanager
+    def zip_opener_helper(self):
+        with zipfile.ZipFile(self.filename) as archive:
+            with archive.open(self.__class__.__qualname__, mode='r') as fd:
+                yield fd
+
+    @contextmanager
+    def zip_closer_helper(self):
+        with zipfile.ZipFile(self.filename, 'w',
+                             compression=zipfile.ZIP_BZIP2) as archive:
+            with io.TextIOWrapper(archive.open(self.__class__.__qualname__,
+                                               mode='w',
+                                               force_zip64=True)) as fd:
+                yield fd
+
     @abc.abstractmethod
     async def load_fd(self, fd):
         pass # pragma: no cover

tests/source/test_file.py

@@ -11,7 +11,7 @@
 import tempfile
 import unittest
 import collections
-from unittest.mock import patch, mock_open
+from unittest.mock import patch, mock_open, Mock
 from functools import wraps
 from contextlib import contextmanager
 from typing import List, Dict, Any, Optional, Tuple, AsyncIterator
@@ -32,10 +32,35 @@ async def repo(self, src_url: str):
         pass # pragma: no cover
 
     async def load_fd(self, fd):
-        pass # pragma: no cover
+        self.loaded_fd = fd
 
     async def dump_fd(self, fd):
-        pass # pragma: no cover
+        self.dumped_fd = fd
+
+class MockZipFile(object):
+
+    def __init__(self, *_args, **_kwargs):
+        self.files = [Mock(filename='foo'), Mock(filename='bar')]
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        return
+
+    def infolist(self):
+        return self.files
+
+    @classmethod
+    def opener(cls):
+        class mzip(cls):
+            pass
+        mzip.open = mock_open()
+        return mzip
+
+@contextmanager
+def yield_42():
+    yield 42
 
 class TestFileSource(AsyncTestCase):
 
@@ -95,6 +120,13 @@ async def test_open_xz(self):
             await source.open()
             m_open.assert_called_once_with('testfile.xz', 'rt')
 
+    async def test_open_zip(self):
+        source = FakeFileSource('testfile.zip')
+        with patch('os.path.exists', return_value=True), \
+                patch.object(source, 'zip_opener_helper', yield_42):
+            await source.open()
+            self.assertEqual(source.loaded_fd, 42)
+
     async def test_open_no_file(self):
         source = FakeFileSource('testfile')
         with patch('os.path.isfile', return_value=False):
@@ -135,7 +167,13 @@ async def test_close_xz(self):
         with patch('lzma.open', m_open):
             await source.close()
             m_open.assert_called_once_with('testfile.xz', 'wt')
-    
+
+    async def test_close_zip(self):
+        source = FakeFileSource('testfile.zip')
+        with patch.object(source, 'zip_closer_helper', yield_42):
+            await source.close()
+            self.assertEqual(source.dumped_fd, 42)
+
     async def test_close_readonly(self):
         source = FakeFileSource('testfile:ro')
         m_open = mock_open()