From 7eeabeccd6693aa8d1ebc42aacced298ed350f3c Mon Sep 17 00:00:00 2001
From: Roel de Jong <12800443+twiggler@users.noreply.github.com>
Date: Wed, 16 Oct 2024 12:22:47 +0200
Subject: [PATCH] Add unit test for tar encryption

---
 tests/conftest.py         |  6 ++++++
 tests/test_outputs_tar.py | 18 ++++++++++++++++++
 tests/test_outputs_zip.py |  8 +-------
 3 files changed, 25 insertions(+), 7 deletions(-)

diff --git a/tests/conftest.py b/tests/conftest.py
index b9d11a8e..7600d629 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -33,3 +33,9 @@ def mock_target(mock_fs: VirtualFilesystem) -> Target:
     target.filesystems.add(mock_fs)
     target.os = "mock"
     return target
+
+
+@pytest.fixture
+def public_key() -> bytes:
+    with open("tests/data/public_key.pem", "r") as f:
+        return f.read()
diff --git a/tests/test_outputs_tar.py b/tests/test_outputs_tar.py
index 0c6e36fd..ad888b4b 100644
--- a/tests/test_outputs_tar.py
+++ b/tests/test_outputs_tar.py
@@ -5,6 +5,7 @@
 from dissect.target.filesystem import VirtualFilesystem
 
 from acquire.outputs import TarOutput
+from acquire.tools.decrypter import EncryptedFile
 
 
 @pytest.fixture(params=[(True, "gzip"), (True, "bzip2"), (True, "xz"), (False, None)])
@@ -41,3 +42,20 @@ def test_tar_output_write_entry(mock_fs: VirtualFilesystem, tar_output: TarOutpu
         assert file.issym()
     elif entry.is_file():
         assert file.isfile()
+
+
+def test_tar_output_encrypt(mock_fs: VirtualFilesystem, public_key: bytes, tmp_path: Path) -> None:
+    entry_name = "/foo/bar/some-file"
+    entry = mock_fs.get(entry_name)
+    tar_output = TarOutput(tmp_path, compress=True, compression_method="gzip", encrypt=True, public_key=public_key)
+    tar_output.write_entry(entry_name, entry)
+    tar_output.close()
+
+    encrypted_stream = EncryptedFile(tar_output.path.open("rb"), Path("tests/data/private_key.pem"))
+    decrypted_path = tmp_path / "decrypted.tar"
+    # Direct streaming is not an otion because tarfile needs seek when reading from encrypted files directly
+    with open(decrypted_path, "wb") as f:
+        f.write(encrypted_stream.read())
+
+    tar_file = tarfile.open(name=decrypted_path, mode="r")
+    assert entry.open().read() == tar_file.extractfile(entry_name).read()
diff --git a/tests/test_outputs_zip.py b/tests/test_outputs_zip.py
index b52613d7..cc556623 100644
--- a/tests/test_outputs_zip.py
+++ b/tests/test_outputs_zip.py
@@ -48,12 +48,6 @@ def test_zip_output_write_entry(mock_fs: VirtualFilesystem, zip_output: ZipOutpu
         assert stat.S_ISREG(file_type)
 
 
-@pytest.fixture
-def public_key() -> bytes:
-    with open("tests/data/public_key.pem", "r") as f:
-        return f.read()
-
-
 def test_zip_output_encrypt(mock_fs: VirtualFilesystem, public_key: bytes, tmp_path: Path) -> None:
     entry_name = "/foo/bar/some-file"
     entry = mock_fs.get(entry_name)
@@ -66,6 +60,6 @@ def test_zip_output_encrypt(mock_fs: VirtualFilesystem, public_key: bytes, tmp_p
     # Direct streaming is not an otion because zipfile needs seek when reading from encrypted files directly
     with open(decrypted_path, "wb") as f:
         f.write(encrypted_stream.read())
-    zip_file = zipfile.ZipFile(decrypted_path, mode="r")
 
+    zip_file = zipfile.ZipFile(decrypted_path, mode="r")
     assert entry.open().read() == zip_file.open(entry_name).read()