Rename pdfraw to pdfany and allow 2 payloads. v1.5.21.5.2

3 files changed, 38 insertions, 23 deletions

diff --git a/PolyglotFile/__init__.py b/PolyglotFile/__init__.py
index 6800028..e6f6416 100644
--- a/PolyglotFile/__init__.py
+++ b/PolyglotFile/__init__.py
@@ -2,7 +2,7 @@
 # -*- coding: utf-8 -*-
 
 from .polyglotpdfzip import PolyglotPdfZip
-from .polyglotpdfraw import PolyglotPdfRaw
+from .polyglotpdfany import PolyglotPdfAny
 from .polyglotzippdf import PolyglotZipPdf
 from .polyglotszippdf import PolyglotSZipPdf
 
diff --git a/PolyglotFile/polyglotpdfraw.py b/PolyglotFile/polyglotpdfany.py
index c524788..8b7289f 100644
--- a/PolyglotFile/polyglotpdfraw.py
+++ b/PolyglotFile/polyglotpdfany.py
@@ -6,10 +6,12 @@ import logging
     |-------------------------------|    -
     |--------- PDF Header ----------K1   | J1
     |-------------------------------|    -
-    |----- PDF OBJ 1 = RAW Data ----K2   |
+    |---- PDF OBJ 1 = Payload 1 ----K2   |
     |-------------------------------|    -
     |---- Original PDF Ojbects -----K3   | J2
     |-------------------------------|    -
+    |-- PDF Last OBJ = Payload 2 ---K4   |
+    |-------------------------------|    |
     |---------- Xref Table ---------|    |
     |-------------------------------K5   |
     |----------- Trailer -----------|    |
@@ -17,21 +19,30 @@ import logging
 '''
 
 
-class PolyglotPdfRaw():
+class PolyglotPdfAny():
     from PdfFileTransformer import Pdf
 
-    def __init__(self, Pdf, Raw_filename):
+    def __init__(self, Pdf, payload1filename, payload2filename):
         self.buffer = bytearray()
         self.pdf = Pdf
-        self.raw_filename = Raw_filename
+        self.payload1 = bytearray()
+        if payload1filename is not None:
+            with open(payload1filename, "rb") as f:
+                self.payload1 = f.read()
+        self.payload2 = bytearray()
+        if payload2filename is not None:
+            with open(payload2filename, "rb") as f:
+                self.payload2 = f.read()
+
         self.buffer = bytearray()
 
     def generate(self):
-        raw_buffer = bytearray()
-        with open(self.raw_filename, "rb") as f:
-            raw_buffer = f.read()
-        k2_stream = raw_buffer
-        self.pdf.insert_new_obj_stream_at_start(k2_stream)
+        k2stream = self.payload1
+        if len(k2stream) > 0:
+            self.pdf.insert_new_obj_stream_at_start(k2stream)
+        k4stream = self.payload2
+        if len(k4stream) > 0:
+            self.pdf.insert_new_obj_stream_at_end(k4stream)
         self.buffer = self.pdf.get_build_buffer()
 
     def write(self, filename):
diff --git a/truepolyglot b/truepolyglot
index 34950ff..de74974 100755
--- a/truepolyglot
+++ b/truepolyglot
@@ -9,7 +9,7 @@ from ZipFileTransformer import Zip
 from PolyglotFile import PolyglotZipPdf
 from PolyglotFile import PolyglotPdfZip
 from PolyglotFile import PolyglotSZipPdf
-from PolyglotFile import PolyglotPdfRaw
+from PolyglotFile import PolyglotPdfAny
 
 
 def main():
@@ -21,10 +21,11 @@ def main():
                        '* szippdf: Generate a file valid as ZIP and PDF.' +
                        ' The format is strictly a ZIP.' +
                        ' Archive is modified.\n' +
-                       '* pdfraw: Generate a file strictly valid as a PDF ' +
-                       'with a custom first object content.')
+                       '* pdfany: Generate a valid PDF file with payload1' +
+                       ' file content as first object or/and payload2 file' +
+                       ' content as last oject.\n')
     usage_str = '%(prog)s format [options] output-file'
-    epilog_str = 'TruePolyglot v1.5.1'
+    epilog_str = 'TruePolyglot v1.5.2'
     frm = argparse.RawTextHelpFormatter
     parser = argparse.ArgumentParser(description=description_str,
                                      epilog=epilog_str,
@@ -33,14 +34,16 @@ def main():
     parser.add_argument('format', nargs='+', choices=["pdfzip",
                                                       "zippdf",
                                                       "szippdf",
-                                                      "pdfraw"],
+                                                      "pdfany"],
                         help='Output polyglot format')
     parser.add_argument('--pdffile', dest='pdffile',
                         help='PDF input file')
     parser.add_argument('--zipfile', dest='zipfile',
                         help='ZIP input file')
-    parser.add_argument('--rawfile', dest='rawfile',
-                        help='RAW input file')
+    parser.add_argument('--payload1file', dest='payload1file',
+                        help='Payload 1 input file')
+    parser.add_argument('--payload2file', dest='payload2file',
+                        help='Payload 2 input file')
     parser.add_argument('--acrobat-compatibility',
                         dest='acrobat_compatibility',
                         help='Add a byte at start for Acrobat Reader compatibility with szippdf format',
@@ -57,11 +60,12 @@ def main():
     if args.acrobat_compatibility and args.format[0] != "szippdf":
         parser.error('--acrobat-compatibility is for szippdf only')
     if "pdf" in args.format[0] and args.pdffile is None:
-            parser.error('pdffile is required')
+        parser.error('pdffile is required')
     if "zip" in args.format[0] and args.zipfile is None:
-            parser.error('zipfile is required')
-    if "raw" in args.format[0] and args.rawfile is None:
-            parser.error('rawfile is required')
+        parser.error('zipfile is required')
+    if ("any" in args.format[0] and args.payload1file is None and
+            args.payload2file is None):
+        parser.error('payload1file or payload2file is required')
 
     if args.verbose == "none":
         logging.basicConfig(level=logging.CRITICAL)
@@ -84,9 +88,9 @@ def main():
         p = Pdf(args.pdffile)
         z = Zip(args.zipfile)
         a = PolyglotSZipPdf(p, z, args.acrobat_compatibility)
-    if args.format[0] == "pdfraw":
+    if args.format[0] == "pdfany":
         p = Pdf(args.pdffile)
-        a = PolyglotPdfRaw(p, args.rawfile)
+        a = PolyglotPdfAny(p, args.payload1file, args.payload2file)
 
     a.generate()
     a.write(args.output_file[0])