Import shared TAR helpers from common.utils

storage_service/common/tests/test_utils.py

@@ -1,5 +1,10 @@
 from __future__ import absolute_import
 from __future__ import unicode_literals
+from collections import namedtuple
+import os
+import shutil
+import subprocess
+import tarfile
 
 from six import StringIO
 import mock
@@ -18,6 +23,9 @@
 COMPRESS_ORDER_ONE = "1"
 COMPRESS_ORDER_TWO = "2"
 
+ExTarCase = namedtuple("ExTarCase", "path isdir raises expected")
+CrTarCase = namedtuple("CrTarCase", "path isfile istar raises expected")
+
 
 @pytest.mark.parametrize(
     "pronom,algorithm,compression",
@@ -290,3 +298,91 @@ def test_package_is_file(package_path, is_file):
     see in the storage service.
     """
     assert utils.package_is_file(package_path) == is_file
+
+
+@pytest.mark.parametrize(
+    "path, will_be_dir, sp_raises, expected",
+    [
+        ExTarCase(path="/a/b/c", isdir=True, raises=False, expected="success"),
+        ExTarCase(path="/a/b/d", isdir=False, raises=True, expected="fail"),
+        ExTarCase(path="/a/b/c", isdir=True, raises=True, expected="fail"),
+    ],
+)
+def test_extract_tar(mocker, path, will_be_dir, sp_raises, expected):
+    if sp_raises:
+        mocker.patch.object(subprocess, "check_output", side_effect=OSError("gotcha!"))
+    else:
+        mocker.patch.object(subprocess, "check_output")
+    mocker.patch.object(os, "rename")
+    mocker.patch.object(os, "remove")
+    if will_be_dir:
+        mocker.patch.object(os.path, "isdir", return_value=True)
+    else:
+        mocker.patch.object(os.path, "isdir", return_value=False)
+    tarpath_ext = "{}.tar".format(path)
+    dirname = os.path.dirname(tarpath_ext)
+    if expected == "success":
+        ret = utils.extract_tar(path)
+        assert ret is None
+        os.remove.assert_called_once_with(tarpath_ext)
+    else:
+        with pytest.raises(utils.TARException) as excinfo:
+            ret = utils.extract_tar(path)
+        assert "Failed to extract {}: gotcha!".format(path) == str(excinfo.value)
+        os.rename.assert_any_call(tarpath_ext, path)
+        assert not os.remove.called
+    os.rename.assert_any_call(path, tarpath_ext)
+    subprocess.check_output.assert_called_once_with(
+        ["tar", "-xf", tarpath_ext, "-C", dirname]
+    )
+
+
+@pytest.mark.parametrize(
+    "path, will_be_file, will_be_tar, sp_raises, expected",
+    [
+        CrTarCase(
+            path="/a/b/c", isfile=True, istar=True, raises=False, expected="success"
+        ),
+        CrTarCase(
+            path="/a/b/c/", isfile=True, istar=True, raises=False, expected="success"
+        ),
+        CrTarCase(
+            path="/a/b/c", isfile=True, istar=False, raises=False, expected="fail"
+        ),
+        CrTarCase(
+            path="/a/b/c", isfile=False, istar=True, raises=False, expected="fail"
+        ),
+        CrTarCase(
+            path="/a/b/c", isfile=False, istar=False, raises=True, expected="fail"
+        ),
+    ],
+)
+def test_create_tar(mocker, path, will_be_file, will_be_tar, sp_raises, expected):
+    if sp_raises:
+        mocker.patch.object(subprocess, "check_output", side_effect=OSError("gotcha!"))
+    else:
+        mocker.patch.object(subprocess, "check_output")
+    mocker.patch.object(os.path, "isfile", return_value=will_be_file)
+    mocker.patch.object(tarfile, "is_tarfile", return_value=will_be_tar)
+    mocker.patch.object(os, "rename")
+    mocker.patch.object(shutil, "rmtree")
+    fixed_path = path.rstrip("/")
+    tarpath = "{}.tar".format(fixed_path)
+    if expected == "success":
+        ret = utils.create_tar(path)
+        shutil.rmtree.assert_called_once_with(fixed_path)
+        os.rename.assert_called_once_with(tarpath, fixed_path)
+        tarfile.is_tarfile.assert_any_call(fixed_path)
+        assert ret is None
+    else:
+        with pytest.raises(utils.TARException) as excinfo:
+            ret = utils.create_tar(path)
+        assert "Failed to create a tarfile at {} for dir at {}".format(
+            tarpath, fixed_path
+        ) == str(excinfo.value)
+        assert not shutil.rmtree.called
+        assert not os.rename.called
+    if not sp_raises:
+        os.path.isfile.assert_called_once_with(tarpath)
+        if will_be_file:
+            tarfile.is_tarfile.assert_any_call(tarpath)

storage_service/common/utils.py

@@ -11,6 +11,7 @@
 import os
 import shutil
 import subprocess
+import tarfile
 import uuid
 
 import scandir
@@ -553,6 +554,80 @@ def set_compression_transforms(aip, compression, transform_order):
     return version, extension, program_name
 
 
+# ########### TAR Packaging ############
+
+
+class TARException(Exception):
+    pass
+
+
+def _abort_create_tar(path, tarpath):
+    fail_msg = _(
+        "Failed to create a tarfile at %(tarpath)s for dir at %(path)s"
+        % {"tarpath": tarpath, "path": path}
+    )
+    LOGGER.error(fail_msg)
+    raise TARException(fail_msg)
+
+
+def create_tar(path):
+    """Create a tarfile from the directory at ``path`` and overwrite
+    ``path`` with that tarfile.
+    """
+    path = path.rstrip("/")
+    tarpath = "{}.tar".format(path)
+    changedir = os.path.dirname(tarpath)
+    source = os.path.basename(path)
+    cmd = ["tar", "-C", changedir, "-cf", tarpath, source]
+    LOGGER.info(
+        "creating archive of %s at %s, relative to %s", source, tarpath, changedir
+    )
+    try:
+        subprocess.check_output(cmd)
+    except (OSError, subprocess.CalledProcessError):
+        _abort_create_tar(path, tarpath)
+
+    # Providing the TAR is successfully created then remove the original.
+    if os.path.isfile(tarpath) and tarfile.is_tarfile(tarpath):
+        try:
+            shutil.rmtree(path)
+        except OSError:
+            # Remove a file-path as We're likely packaging a file, e.g. 7z.
+            os.remove(path)
+        os.rename(tarpath, path)
+    else:
+        _abort_create_tar(path, tarpath)
+    try:
+        assert tarfile.is_tarfile(path)
+        assert not os.path.exists(tarpath)
+    except AssertionError:
+        _abort_create_tar(path, tarpath)
+
+
+def _abort_extract_tar(tarpath, newtarpath, err):
+    fail_msg = _(
+        "Failed to extract %(tarpath)s: %(error)s" % {"tarpath": tarpath, "error": err}
+    )
+    LOGGER.error(fail_msg)
+    os.rename(newtarpath, tarpath)
+    raise TARException(fail_msg)
+
+
+def extract_tar(tarpath):
+    """Extract tarfile at ``path`` to a directory at ``path``."""
+    newtarpath = "{}.tar".format(tarpath)
+    os.rename(tarpath, newtarpath)
+    changedir = os.path.dirname(newtarpath)
+    cmd = ["tar", "-xf", newtarpath, "-C", changedir]
+    try:
+        subprocess.check_output(cmd)
+    except (OSError, subprocess.CalledProcessError) as err:
+        _abort_extract_tar(tarpath, newtarpath, err)
+    # TODO: GPG treats this differently because it only ever expects to
+    # TAR a directory but we actually want to TAR file-types as well.
+    os.remove(newtarpath)
+
+
 # ########### OTHER ############
 
 

storage_service/locations/models/gpg.py

@@ -4,8 +4,6 @@
 import datetime
 import logging
 import os
-import shutil
-import subprocess
 import tarfile
 
 # Core Django, alphabetical
@@ -242,7 +240,7 @@ def _gpg_encrypt(path, key_fingerprint):
     """
     tar_created = False
     if os.path.isdir(path):
-        _create_tar(path)
+        utils.create_tar(path)
         tar_created = True
     encr_path, result = gpgutils.gpg_encrypt_file(path, key_fingerprint)
     if os.path.isfile(encr_path) and result.ok:
@@ -252,7 +250,7 @@ def _gpg_encrypt(path, key_fingerprint):
         return path, result
     else:
         if tar_created:
-            _extract_tar(path)
+            utils.extract_tar(path)
         fail_msg = _(
             "An error occured when attempting to encrypt" " %(path)s" % {"path": path}
         )
@@ -291,69 +289,6 @@ def _encr_path2key_fingerprint(encr_path):
         raise GPGException(fail_msg)
 
 
-def _abort_create_tar(path, tarpath):
-    fail_msg = _(
-        "Failed to create a tarfile at %(tarpath)s for dir at %(path)s"
-        % {"tarpath": tarpath, "path": path}
-    )
-    LOGGER.error(fail_msg)
-    raise GPGException(fail_msg)
-
-
-def _create_tar(path):
-    """Create a tarfile from the directory at ``path`` and overwrite ``path``
-    with that tarfile.
-    """
-    path = path.rstrip("/")
-    tarpath = "{}.tar".format(path)
-    changedir = os.path.dirname(tarpath)
-    source = os.path.basename(path)
-    cmd = ["tar", "-C", changedir, "-cf", tarpath, source]
-    LOGGER.info(
-        "creating archive of %s at %s, relative to %s", source, tarpath, changedir
-    )
-    try:
-        subprocess.check_output(cmd)
-    except (OSError, subprocess.CalledProcessError):
-        _abort_create_tar(path, tarpath)
-    if os.path.isfile(tarpath) and tarfile.is_tarfile(tarpath):
-        shutil.rmtree(path)
-        os.rename(tarpath, path)
-    else:
-        _abort_create_tar(path, tarpath)
-    try:
-        assert tarfile.is_tarfile(path)
-        assert not os.path.exists(tarpath)
-    except AssertionError:
-        _abort_create_tar(path, tarpath)
-
-
-def _abort_extract_tar(tarpath, newtarpath):
-    fail_msg = _(
-        "Failed to extract %(tarpath)s to a directory at the same"
-        " location." % {"tarpath": tarpath}
-    )
-    LOGGER.error(fail_msg)
-    os.rename(newtarpath, tarpath)
-    raise GPGException(fail_msg)
-
-
-def _extract_tar(tarpath):
-    """Extract tarfile at ``path`` to a directory at ``path``."""
-    newtarpath = "{}.tar".format(tarpath)
-    os.rename(tarpath, newtarpath)
-    changedir = os.path.dirname(newtarpath)
-    cmd = ["tar", "-xf", newtarpath, "-C", changedir]
-    try:
-        subprocess.check_output(cmd)
-    except (OSError, subprocess.CalledProcessError):
-        _abort_extract_tar(tarpath, newtarpath)
-    if os.path.isdir(tarpath):
-        os.remove(newtarpath)
-    else:
-        _abort_extract_tar(tarpath, newtarpath)
-
-
 def _parse_gpg_version(raw_gpg_version):
     return ".".join(str(i) for i in raw_gpg_version)
 
@@ -392,7 +327,7 @@ def _gpg_decrypt(path):
     # using an uncompressed AIP as input. We extract those here.
     if tarfile.is_tarfile(path) and os.path.splitext(path)[1] == "":
         LOGGER.info("%s is a tarfile so we are extracting it", path)
-        _extract_tar(path)
+        utils.extract_tar(path)
     return path