From a3cb5f13edf5fa0bf69689e8d220831bdea49531 Mon Sep 17 00:00:00 2001
From: Nich <onicrypt@gmail.com>
Date: Mon, 28 Mar 2016 06:04:27 +0000
Subject: [PATCH] Fixed bugs with crypto.py, need to work on cleanup

Project is now properly encrypting the archive,
and now I just need to ensure proper decryption
is happening. Also need to work on cleaning up
clear text versions of the crypt project files.
Need to write tests for flushing out edge cases.
---
 pappyproxy/compress.py | 38 +++++++++++++++++---------------------
 pappyproxy/crypto.py   |  6 +++---
 pappyproxy/pappy.py    |  9 +++++++--
 3 files changed, 27 insertions(+), 26 deletions(-)

diff --git a/pappyproxy/compress.py b/pappyproxy/compress.py
index 662f30b..28d1e5d 100644
--- a/pappyproxy/compress.py
+++ b/pappyproxy/compress.py
@@ -23,19 +23,19 @@ class Compress(object):
         self.zip_archive = sessconfig.archive
         self.bz2_archive = sessconfig.archive
 
-    def compress_project():
+    def compress_project(self):
         if bz2:
-            tar_project()
+            self.tar_project()
         else:
-            zip_project()
+            self.zip_project()
     
-    def decompress_project():
+    def decompress_project(self):
         if bz2:
-            untar_project()
+            self.untar_project()
         else:
-            unzip_project()
+            self.unzip_project()
     
-    def zip_project():
+    def zip_project(self):
         """
         Zip project files
     
@@ -52,7 +52,7 @@ class Compress(object):
             raise PappyException("Error creating the zipfile", e)
         pass
     
-    def unzip_project():
+    def unzip_project(self):
         """
         Extract project files from decrypted zip archive.
         Initially checks the zip archive's magic number and
@@ -71,24 +71,20 @@ class Compress(object):
     
         zf.extractall()
     
-    def tar_project():
-        if tarfile.is_tarfile(self.bz2_archive):
-            archive = tarfile.open(self.bz2_archive, 'w:bz2')
-            project_files = self.config.get_project_files()
-    
-            # Read files line by line to accomodate larger files, e.g. the project database
-            for pf in project_files:
-                archive.add(pf)
+    def tar_project(self):
+        archive = tarfile.open(self.bz2_archive, 'w:bz2')
+        project_files = self.config.get_project_files()
+         
+        # Read files line by line to accomodate larger files, e.g. the project database
+        for pf in project_files:
+            archive.add(pf)
     	archive.close()
     
-    def untar_project():
+    def untar_project(self):
         if tarfile.is_tarfile(self.bz2_archive):
-            # Attempt to read the first 16 bytes of the archive
             # Raise exception if there is a failure
-    	    project_files = self.config.get_project_files()
             try:
                 with tarfile.open(self.bz2_archive, "r:bz2") as archive:
-                    for pf in project_files:
-                        archive.add(pf)
+                    archive.extractall()
             except e:
                 raise PappyException("Project archive contents corrupted. Error: ", e)
diff --git a/pappyproxy/crypto.py b/pappyproxy/crypto.py
index 9194e1b..d8fa3d2 100644
--- a/pappyproxy/crypto.py
+++ b/pappyproxy/crypto.py
@@ -36,7 +36,7 @@ class Crypto(object):
    
         # Create project archive and crypto archive 
         self.compressor.compress_project()
-        archive_file = open(self.archive, 'rb') 
+        archive_file = open(self.archive, 'rb').read() 
         archive_crypt = open(self.config.crypt_file, 'wb')
     
         # Encrypt the archive read as a bytestring
@@ -57,7 +57,7 @@ class Crypto(object):
         # Get the password and salt, then derive the key
         self.crypto_ramp_up()
 
-        crypto_path = os.path.join(os.getcwd(), self.config.crypt_dir)
+        crypto_path = self.config.crypt_dir
         
         if not os.path.isdir(crypto_path):
             os.mkdir(crypto_path)
@@ -168,7 +168,7 @@ class Crypto(object):
     
         try:
             if not self.key:
-                self.key = scrypt.hash(self.password, self.salt, buflen=32)
+                self.key = b64encode(scrypt.hash(self.password, self.salt, buflen=32))
         except TypeError, e:
             raise PappyException("Scrypt failed with type error: ", e)
         except scrypt.error, e:
diff --git a/pappyproxy/pappy.py b/pappyproxy/pappy.py
index fa03160..5f92190 100755
--- a/pappyproxy/pappy.py
+++ b/pappyproxy/pappy.py
@@ -145,11 +145,11 @@ class PappySession(object):
 
     @defer.inlineCallbacks
     def encrypt(self):
-        self.crypto.encrypt_project(self.password)     
+        yield self.crypto.encrypt_project()     
 
     @defer.inlineCallbacks
     def decrypt(self):
-        self.crypto.decrypt_project()
+        yield self.crypto.decrypt_project()
                 
     @defer.inlineCallbacks
     def cleanup(self, ignored=None):
@@ -159,6 +159,11 @@ class PappySession(object):
         if self.delete_data_on_quit:
             print 'Deleting temporary datafile'
             os.remove(self.config.datafile)
+
+        # If currently in the crypt directory,
+        # encrypt the project, delete clear files
+        if os.getcwd() == self.config.crypt_dir:
+            self.encrypt()
     
 def parse_args():
     # parses sys.argv and returns a settings dictionary