Add PGP metadata extraction and tests

Implemented a get_metadata() method in the PGP unpacker to identify PGP file types (public key, private key, encrypted message, signature) from both ASCII-armored and binary formats. Added unit tests to verify metadata extraction for various PGP data types. Improved resource cleanup in the unpack method.

Author: doomedraven

sflock/unpack/pgp.py

@@ -4,6 +4,7 @@
 # See the file 'docs/LICENSE.txt' for copying permission.
 
 import os
+import shutil
 import subprocess
 import tempfile
 
@@ -38,7 +39,7 @@ def unpack(self, password: str = None, duplicates=None):
                 stderr=subprocess.PIPE,
             )
 
-            stdout, stderr = p.communicate(timeout=30)
+            _, _ = p.communicate(timeout=30)
             return_code = p.returncode
 
         except subprocess.TimeoutExpired:
@@ -51,12 +52,45 @@ def unpack(self, password: str = None, duplicates=None):
                 p.kill()
                 p.wait()
             return_code = 1
+        finally:
+            if temporary and os.path.exists(filepath):
+                os.unlink(filepath)
+
+            if os.path.exists(dirpath):
+                shutil.rmtree(dirpath)
+
 
         ret = not return_code
         if not ret:
             return []
 
-        if temporary:
-            os.unlink(filepath)
+    def get_metadata(self):
+        ret = []
+        content = self.f.contents
+        if not content:
+            return ret
+
+        if b"BEGIN PGP PUBLIC KEY BLOCK" in content:
+            ret.append("public_key")
+        elif b"BEGIN PGP PRIVATE KEY BLOCK" in content:
+            ret.append("private_key")
+        elif b"BEGIN PGP MESSAGE" in content:
+            ret.append("encrypted_message")
+        elif b"BEGIN PGP SIGNATURE" in content:
+            ret.append("signature")
+        elif content[0] & 0x80:
+            # Binary analysis
+            tag = content[0]
+            if tag & 0x40:  # New format
+                tag_type = tag & 0x3F
+            else:  # Old format
+                tag_type = (tag >> 2) & 0xF
+
+            if tag_type in (6, 14):
+                ret.append("public_key")
+            elif tag_type == 5:
+                ret.append("private_key")
+            elif tag_type in (1, 18):
+                ret.append("encrypted_message")
 
-        return self.process_directory(dirpath, duplicates, password)
+        return ret

tests/test_pgp_metadata.py

@@ -0,0 +1,40 @@
+import unittest
+from sflock.unpack.pgp import PGP
+from sflock.abstracts import File
+
+class TestPGPMetadata(unittest.TestCase):
+    def test_public_key_ascii(self):
+        f = File(contents=b"-----BEGIN PGP PUBLIC KEY BLOCK-----...")
+        p = PGP(f)
+        self.assertEqual(p.get_metadata(), ["public_key"])
+
+    def test_private_key_ascii(self):
+        f = File(contents=b"-----BEGIN PGP PRIVATE KEY BLOCK-----...")
+        p = PGP(f)
+        self.assertEqual(p.get_metadata(), ["private_key"])
+
+    def test_encrypted_ascii(self):
+        f = File(contents=b"-----BEGIN PGP MESSAGE-----...")
+        p = PGP(f)
+        self.assertEqual(p.get_metadata(), ["encrypted_message"])
+
+    def test_binary_public_key_old(self):
+        # Tag 6 (Public Key Packet) -> 000110
+        # Old format: 10xxxxxx
+        # Tag is bits 5-2.
+        # 0x80 | (6 << 2) = 0x80 | 0x18 = 0x98.
+        # Let's use 0x98 (length type 0 - 1 byte length)
+        f = File(contents=b"\x98\x01") 
+        p = PGP(f)
+        self.assertEqual(p.get_metadata(), ["public_key"])
+
+    def test_binary_encrypted_new(self):
+        # Tag 18 (Sym Encrypted) -> 010010
+        # New format: 11xxxxxx
+        # 0xC0 | 18 = 0xC0 | 0x12 = 0xD2
+        f = File(contents=b"\xD2\x05")
+        p = PGP(f)
+        self.assertEqual(p.get_metadata(), ["encrypted_message"])
+
+if __name__ == '__main__':
+    unittest.main()