Allow customization of ASCII armor comment and version headers

pgpainless-core/src/main/java/org/pgpainless/util/ArmoredOutputStreamFactory.java

@@ -24,11 +24,62 @@ import org.bouncycastle.bcpg.ArmoredOutputStream;
  */
 public class ArmoredOutputStreamFactory {
 
-    public static final String VERSION = "PGPainless";
+    public static final String PGPAINLESS = "PGPainless";
+    private static String VERSION = PGPAINLESS;
+    public static String[] COMMENT = new String[0];
 
     public static ArmoredOutputStream get(OutputStream outputStream) {
         ArmoredOutputStream armoredOutputStream = new ArmoredOutputStream(outputStream);
         armoredOutputStream.setHeader(ArmorUtils.HEADER_VERSION, VERSION);
+        for (String comment : COMMENT) {
+            ArmorUtils.addCommentHeader(armoredOutputStream, comment);
+        }
         return armoredOutputStream;
     }
+
+    /**
+     * Overwrite the version header of ASCII armors with a custom value.
+     * Newlines in the version info string result in multiple version header entries.
+     *
+     * @param version version string
+     */
+    public static void setVersionInfo(String version) {
+        if (version == null || version.trim().isEmpty()) {
+            throw new IllegalArgumentException("Version Info MUST NOT be null NOR empty.");
+        }
+        VERSION = version;
+    }
+
+    /**
+     * Reset the version header to its default value of {@link #PGPAINLESS}.
+     */
+    public static void resetVersionInfo() {
+        VERSION = PGPAINLESS;
+    }
+
+    /**
+     * Set a comment header value in the ASCII armor header.
+     * If the comment contains newlines, it will be split into multiple header entries.
+     *
+     * @param comment comment
+     */
+    public static void setComment(String comment) {
+        if (comment == null) {
+            throw new IllegalArgumentException("Comment cannot be null.");
+        }
+        String trimmed = comment.trim();
+        if (trimmed.isEmpty()) {
+            throw new IllegalArgumentException("Comment cannot be empty.");
+        }
+
+        String[] lines = comment.split("\n");
+        COMMENT = lines;
+    }
+
+    /**
+     * Reset to the default of no comment headers.
+     */
+    public static void resetComment() {
+        COMMENT = new String[0];
+    }
 }

pgpainless-core/src/test/java/org/pgpainless/util/ArmorUtilsTest.java

@@ -22,13 +22,21 @@ import static org.junit.jupiter.api.Assertions.assertTrue;
 import java.io.ByteArrayInputStream;
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
+import java.io.InputStream;
 import java.nio.charset.StandardCharsets;
 import java.util.List;
 
 import org.bouncycastle.bcpg.ArmoredInputStream;
 import org.bouncycastle.bcpg.ArmoredOutputStream;
+import org.bouncycastle.openpgp.PGPCompressedData;
 import org.bouncycastle.openpgp.PGPException;
+import org.bouncycastle.openpgp.PGPLiteralData;
+import org.bouncycastle.openpgp.PGPObjectFactory;
 import org.bouncycastle.openpgp.PGPSecretKeyRing;
+import org.bouncycastle.openpgp.PGPUtil;
+import org.bouncycastle.openpgp.bc.BcPGPObjectFactory;
+import org.bouncycastle.util.io.Streams;
+import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.Test;
 import org.pgpainless.algorithm.HashAlgorithm;
 import org.pgpainless.key.TestKeys;
@@ -120,4 +128,54 @@ public class ArmorUtilsTest {
         String ascii = ArmorUtils.toAsciiArmoredString(in);
         assertTrue(ascii.startsWith("-----BEGIN PGP PRIVATE KEY BLOCK-----\n"));
     }
+
+    @Test
+    public void testSetCustomVersionHeader() throws IOException {
+        ArmoredOutputStreamFactory.setVersionInfo("MyVeryFirstOpenPGPProgram 1.0");
+        ArmoredOutputStreamFactory.setComment("This is a comment\nThat spans multiple\nLines!");
+
+        ByteArrayOutputStream out = new ByteArrayOutputStream();
+        ArmoredOutputStream armorOut = ArmoredOutputStreamFactory.get(out);
+
+        byte[] data = "This is a very secret message that nobody is allowed to read.".getBytes(StandardCharsets.UTF_8);
+        armorOut.write(data);
+        armorOut.close();
+
+        assertEquals("-----BEGIN PGP MESSAGE-----\n" +
+                "Version: MyVeryFirstOpenPGPProgram 1.0\n" +
+                "Comment: This is a comment\n" +
+                "Comment: That spans multiple\n" +
+                "Comment: Lines!\n" +
+                "\n" +
+                "VGhpcyBpcyBhIHZlcnkgc2VjcmV0IG1lc3NhZ2UgdGhhdCBub2JvZHkgaXMgYWxs\n" +
+                "b3dlZCB0byByZWFkLg==\n" +
+                "=XMZb\n" +
+                "-----END PGP MESSAGE-----\n", out.toString());
+    }
+
+    @Test
+    public void decodeExampleTest() throws IOException, PGPException {
+        String armored = "-----BEGIN PGP MESSAGE-----\n" +
+                "Version: OpenPrivacy 0.99\n" +
+                "\n" +
+                "yDgBO22WxBHv7O8X7O/jygAEzol56iUKiXmV+XmpCtmpqQUKiQrFqclFqUDBovzS\n" +
+                "vBSFjNSiVHsuAA==\n" +
+                "=njUN\n" +
+                "-----END PGP MESSAGE-----";
+        InputStream inputStream = PGPUtil.getDecoderStream(new ByteArrayInputStream(armored.getBytes(StandardCharsets.UTF_8)));
+        PGPObjectFactory factory = new BcPGPObjectFactory(inputStream);
+        PGPCompressedData compressed = (PGPCompressedData) factory.nextObject();
+        factory = new BcPGPObjectFactory(compressed.getDataStream());
+        PGPLiteralData literal = (PGPLiteralData) factory.nextObject();
+        ByteArrayOutputStream out = new ByteArrayOutputStream();
+        assertEquals("_CONSOLE", literal.getFileName());
+        Streams.pipeAll(literal.getInputStream(), out);
+        assertEquals("Can't anyone keep a secret around here?\n", out.toString());
+    }
+
+    @AfterAll
+    public static void resetHeaders() {
+        ArmoredOutputStreamFactory.resetComment();
+        ArmoredOutputStreamFactory.resetVersionInfo();
+    }
 }