From d11beb10af043f77df8c6fa6c0e3150c642ad158 Mon Sep 17 00:00:00 2001
From: r4sas <r4sas@i2pmail.org>
Date: Fri, 7 Jun 2019 09:12:00 +0000
Subject: [PATCH] Make compression optional

ref: https://github.com/PrivateBin/PrivateBin/issues/38
Compression is optional only in v2 paste format.
Currently available types: 'zlib' and 'none'.
To set compression use `-c "none"` or `--compression "none".

When receiving paste, compression detected from 'adata' field
---
 pbincli/actions.py | 11 ++++++++++-
 pbincli/cli.py     |  2 ++
 pbincli/format.py  | 43 ++++++++++++++++++++++++++++++++-----------
 3 files changed, 44 insertions(+), 12 deletions(-)

diff --git a/pbincli/actions.py b/pbincli/actions.py
index b56ee7e..3df38d8 100644
--- a/pbincli/actions.py
+++ b/pbincli/actions.py
@@ -14,7 +14,16 @@ def send(args, api_client):
         text = ""
 
     paste = Paste(args.debug)
-    paste.setVersion(api_client.getVersion())
+
+    # get from server supported paste format version and update object
+    version = api_client.getVersion()
+    paste.setVersion(version)
+
+    # set compression type, works only on v2 pastes
+    if version == 2:
+        paste.setCompression(args.compression)
+
+    # add text in paste (if it provided)
     paste.setText(text)
 
     # If we set PASSWORD variable
diff --git a/pbincli/cli.py b/pbincli/cli.py
index bbd3c39..d4356df 100755
--- a/pbincli/cli.py
+++ b/pbincli/cli.py
@@ -36,6 +36,8 @@ def main():
     send_parser.add_argument("-t", "--text", help="comment in quotes. Ignored if used stdin")
     send_parser.add_argument("-q", "--notext", default=False, action="store_true", help="don't send text in paste")
     send_parser.add_argument("-p", "--password", help="password for encrypting paste")
+    send_parser.add_argument("-c", "--compression", default="zlib", action="store",
+        choices=["zlib", "none"], help="set compression for paste (default: zlib). Note: works only on v2 paste format")
     send_parser.add_argument("-d", "--debug", default=False, action="store_true", help="enable debug")
     send_parser.add_argument("--dry", default=False, action="store_true", help="invoke dry run")
     send_parser.add_argument("-f", "--file", help="example: image.jpg or full path to file")
diff --git a/pbincli/format.py b/pbincli/format.py
index 6e37f21..7edf01e 100644
--- a/pbincli/format.py
+++ b/pbincli/format.py
@@ -1,6 +1,7 @@
 from Crypto.Random import get_random_bytes
 from Crypto.Cipher import AES
 from base64 import b64encode, b64decode
+from pbincli.utils import PBinCLIException
 
 CIPHER_ITERATION_COUNT = 100000
 CIPHER_SALT_BYTES = 8
@@ -12,12 +13,13 @@ CIPHER_TAG_BYTES = int(CIPHER_TAG_BITS/8)
 class Paste:
     def __init__(self, debug=False):
         self._version = 2
-        self._data = ""
-        self._text = ""
-        self._attachment = ""
-        self._attachment_name = ""
+        self._compression = 'zlib'
+        self._data = ''
+        self._text = ''
+        self._attachment = ''
+        self._attachment_name = ''
         self._key = get_random_bytes(CIPHER_BLOCK_BYTES)
-        self._password = ""
+        self._password = ''
         self._debug = debug
 
 
@@ -52,6 +54,9 @@ class Paste:
         self._attachment = 'data:' + mime + ';base64,' + b64encode(contents).decode()
         self._attachment_name = path_leaf(path)
 
+    def setCompression(self, comp):
+        self._compression = comp
+
 
     def getText(self):
         return self._text
@@ -117,7 +122,14 @@ class Paste:
     def __decompress(self, s):
         import zlib
         if self._version == 2:
-            return zlib.decompress(s, -zlib.MAX_WBITS)
+            if self._compression == 'zlib':
+                # decompress data
+                return zlib.decompress(s, -zlib.MAX_WBITS)
+            elif self._compression == 'none':
+                # nothing to do, just return original data
+                return s
+            else:
+                raise PBinCLIException('Unknown compression type provided in paste!')
         else:
             return zlib.decompress(bytearray(map(ord, b64decode(s.encode('utf-8')).decode('utf-8'))), -zlib.MAX_WBITS)
 
@@ -125,10 +137,16 @@ class Paste:
     def __compress(self, s):
         import zlib
         if self._version == 2:
-            # using compressobj as compress doesn't let us specify wbits
-            # needed to get the raw stream without headers
-            co = zlib.compressobj(wbits=-zlib.MAX_WBITS)
-            return co.compress(s) + co.flush()
+            if self._compression == 'zlib':
+                # using compressobj as compress doesn't let us specify wbits
+                # needed to get the raw stream without headers
+                co = zlib.compressobj(wbits=-zlib.MAX_WBITS)
+                return co.compress(s) + co.flush()
+            elif self._compression == 'none':
+                # nothing to do, just return original data
+                return s
+            else:
+                raise PBinCLIException('Unknown compression type provided!')
         else:
             co = zlib.compressobj(wbits=-zlib.MAX_WBITS)
             b = co.compress(s) + co.flush()
@@ -143,6 +161,9 @@ class Paste:
             salt = b64decode(self._data['adata'][0][1])
             key = self.__deriveKey(salt)
 
+            # Get compression type from received paste
+            self._compression = self._data['adata'][0][7]
+
             cipher = self.__initializeCipher(key, iv, self._data['adata'])
             # Cut the cipher text into message and tag
             cipher_text_tag = b64decode(self._data['ct'])
@@ -203,7 +224,7 @@ class Paste:
                     CIPHER_TAG_BITS,
                     'aes',
                     'gcm',
-                    'zlib'
+                    self._compression
                 ],
                 formatter,
                 int(burnafterreading),